package gold.contest;

import gold.utils.InputUtil;

import java.math.BigDecimal;
import java.util.*;
import java.util.stream.Collectors;

/**
 * Created by fanzhenyu02 on 2021/12/10.
 * common problem solver template.
 */
public class LC6121 {
    public long startExecuteTime = System.currentTimeMillis();


    class Solution {
        public int[] smallestTrimmedNumbers(String[] nums, int[][] queries) {
            int[] ans = new int[queries.length];
            Map<Integer, TreeMap<String, TreeSet<Integer>>> trimMap = new HashMap<>(); // trim -> num - > id List
            Map<String, Integer> memo = new HashMap<>();

            for (int i = 0; i < queries.length; i++) {
                int trim = queries[i][1], k = queries[i][0];
                String queryKey = trim + "_" + k;
                if (memo.containsKey(queryKey)) ans[i] = memo.get(queryKey);
                if (!trimMap.containsKey(trim)) {
                    addToTrim(nums, trim, trimMap);
                }

                final TreeMap<String, TreeSet<Integer>> bigDecimalTreeSetTreeMap = trimMap.get(trim);
                int cntK = 0;
                for (String curNum : bigDecimalTreeSetTreeMap.keySet()) {
                    if (cntK + bigDecimalTreeSetTreeMap.get(curNum).size() >= k) {
                        ans[i] = bigDecimalTreeSetTreeMap.get(curNum).stream().collect(Collectors.toList()).get(k - cntK - 1);
                        memo.put(queryKey, ans[i]);
                        break;
                    } else {
                        cntK += bigDecimalTreeSetTreeMap.get(curNum).size();
                    }
                }
            }

            return ans;
        }

        public void addToTrim(String[] nums, int trim, Map<Integer, TreeMap<String, TreeSet<Integer>>> trimMap) {
            trimMap.put(trim, new TreeMap<>());
            for (int i = 0; i < nums.length; i++) {
                String num = nums[i];
                String curNum = num.substring(num.length() - trim, num.length());
                if (!trimMap.get(trim).containsKey(curNum)) trimMap.get(trim).put(curNum, new TreeSet<>());
                trimMap.get(trim).get(curNum).add(i);
            }
        }
    }


    class Solution_BigDecimal_Version {
        public int[] smallestTrimmedNumbers(String[] nums, int[][] queries) {
            int[] ans = new int[queries.length];
            Map<Integer, TreeMap<BigDecimal, TreeSet<Integer>>> trimMap = new HashMap<>(); // trim -> num - > id List
            Map<String, Integer> memo = new HashMap<>();

            for (int i = 0; i < queries.length; i++) {
                int trim = queries[i][1], k = queries[i][0];
                String queryKey = trim + "_" + k;
                if (memo.containsKey(queryKey)) ans[i] = memo.get(queryKey);
                if (!trimMap.containsKey(trim)) {
                    addToTrim(nums, trim, trimMap);
                }

                final TreeMap<BigDecimal, TreeSet<Integer>> bigDecimalTreeSetTreeMap = trimMap.get(trim);
                int cntK = 0;
                for (BigDecimal curNum : bigDecimalTreeSetTreeMap.keySet()) {
                    if (cntK + bigDecimalTreeSetTreeMap.get(curNum).size() >= k) {
                        ans[i] = bigDecimalTreeSetTreeMap.get(curNum).stream().collect(Collectors.toList()).get(k - cntK - 1);
                        memo.put(queryKey, ans[i]);
                        break;
                    } else {
                        cntK += bigDecimalTreeSetTreeMap.get(curNum).size();
                    }
                }
            }

            return ans;
        }

        public void addToTrim(String[] nums, int trim, Map<Integer, TreeMap<BigDecimal, TreeSet<Integer>>> trimMap) {
            trimMap.put(trim, new TreeMap<>());
            for (int i = 0; i < nums.length; i++) {
                String num = nums[i];
                BigDecimal curNum = new BigDecimal(num.substring(num.length() - trim, num.length()));
                if (!trimMap.get(trim).containsKey(curNum)) trimMap.get(trim).put(curNum, new TreeSet<>());
                trimMap.get(trim).get(curNum).add(i);
            }
        }
    }

    public void run() {
        System.out.println(Arrays.toString(new Solution().smallestTrimmedNumbers(new String[]{"102", "473", "251", "814"}, InputUtil.toDoubleIntegerArray("[[1,1],[2,3],[4,2],[1,2]]"))));
        System.out.println(Arrays.toString(new Solution().smallestTrimmedNumbers(new String[]{"24", "37", "96", "04"}, InputUtil.toDoubleIntegerArray("[[2,1],[2,2]]"))));

    }

    public static void main(String[] args) throws Exception {
        LC6121 an = new LC6121();
        an.run();

        System.out.println("\ncurrent solution total execute time: " + (System.currentTimeMillis() - an.startExecuteTime) + " ms.");
    }
}
